//
// Created by Luhang on 2023/5/25/025.
//
#pragma once


#include <iostream>
#include <ctime>
using namespace std;

enum Colour{
    RED,
    BLACK,
};

template<class T>
struct RBTreeNode{
    RBTreeNode<T>* _left;
    RBTreeNode<T>* _right;
    RBTreeNode<T>* _parent;
    T _data;
    Colour _col;

    RBTreeNode(const T& data)
            :_left(nullptr)
            ,_right(nullptr)
            ,_parent(nullptr)
            ,_data(data)
            ,_col(RED)
    {}
};

template<class T, class Ref, class Ptr>
struct __RBTreeIterator{
    typedef RBTreeNode<T> Node;
    typedef __RBTreeIterator<T, Ref, Ptr> Self;
    Node* _node;

    __RBTreeIterator(Node* node)
    :_node(node)
    {}

    //无论__RBTreeIterator被实例化成什么迭代器，__RBTreeIterator<T, T&, T*>都是个普通迭代器
    //即用普通迭代器去构造普通迭代器（拷贝构造）或初始化const迭代器（构造函数）
    __RBTreeIterator(const __RBTreeIterator<T, T&, T*>&  it)
    :_node(it._node)
    {}

    Ref operator*(){
        return _node->_data;
    }

    Ptr operator->(){
        return &_node->_data;
    }

    bool operator!=(const Self& s){
        return _node != s._node;
    }

    Self& operator++(){
        //右不为空，则下一个就是右子树的最左节点
        if (_node->_right){
            Node* subLeft = _node->_right;
            while (subLeft->_left){
                subLeft = subLeft->_left;
            }
            _node = subLeft;
        }
        //右为空，要么是子树遍历完毕该上层子树的根，要么是作为子树的左节点该子树的根
        //总之沿着根的路径，要找子为父的左的那个祖先
        else{
            Node* cur = _node;
            Node* parent = _node->_parent;
            while (parent && cur == parent->_right){
                cur = parent;
                parent = cur->_parent;
            }
            _node = parent;
        }
        return *this;
    }

    Self& operator--(){
        //左不为空，则下一个就是左子树的最右节点
        if (_node->_left){
            Node* subRight = _node->_left;
            while (subRight->_right){
                subRight = subRight->_right;
            }
            _node = subRight;
        }
        //左为空，找子为父的右的那个祖先
        else{
            Node* cur = _node;
            Node* parent = _node->_parent;
            while (parent && cur == parent->_left){
                cur = parent;
                parent = cur->_parent;
            }
            _node = parent;
        }
        return *this;
    }

};


template<class K, class T, class KeyOfT>
class RBTree{
    typedef RBTreeNode<T> Node;
public:
    ~RBTree(){
        _Destroy(_root);
        _root = nullptr;
    }

    typedef __RBTreeIterator<T, T&, T*> iterator;
    typedef __RBTreeIterator<T, const T&, const T*> const_iterator;

    iterator begin(){
        Node* cur = _root;
        while (cur && cur->_left)
            cur = cur->_left;
        return iterator(cur);
    }

    iterator end(){
        return iterator(nullptr);
    }

    const_iterator begin() const{
        Node* cur = _root;
        while (cur && cur->_left)
            cur = cur->_left;
        return const_iterator(cur);
    }

    const_iterator end() const{
        return const_iterator(nullptr);
    }
    
    Node* Find(const K& key){
        Node* cur = _root;
        KeyOfT kot;
        while (cur){
            if (kot(cur->_data) < key)
                cur = cur->_right;
            else if (kot(cur->_data) > key)
                cur = cur->_left;
            else
                return cur;
        }
        return nullptr;
    }

    pair<iterator,bool> Insert(const T& data){
        if (_root == nullptr){
            _root = new Node(data);
            _root->_col = BLACK;
            return make_pair(iterator(_root), true);
        }

        KeyOfT kot;
        Node* parent = nullptr;
        Node* cur = _root;
        while (cur){
            if (kot(cur->_data) < kot(data)){
                parent = cur;
                cur = cur->_right;
            }
            else if (kot(cur->_data) > kot(data)){
                parent = cur;
                cur = cur->_left;
            } else
                return make_pair(iterator(cur), false);
        }

        //链接
        cur = new Node(data);
        Node* newNode = cur;

        if (kot(parent->_data) > kot(data))
            parent->_left = cur;
        else
            parent->_right = cur;
        cur->_parent = parent;

        //开始修正
        while (parent && parent->_col == RED){
            Node* grandfather = parent->_parent;
            //p在g的左
            if (grandfather->_left == parent){
                Node* uncle = grandfather->_right;
                //情况1：u存在且为红(变色处理，并继续往上处理)
                if (uncle && uncle->_col == RED){
                    parent->_col = BLACK;
                    uncle->_col = BLACK;
                    grandfather->_col = RED;

                    //继续向上调整
                    cur = grandfather;
                    parent = cur->_parent;
                }
                    //情况2：u不存在/u存在且为黑(旋转+变色)
                else{
                    //斜线:单旋
                    //    g
                    //  p   u
                    //c
                    if (cur == parent->_left){
                        RotateR(grandfather);
                        parent->_col = BLACK;
                        grandfather->_col = RED;
                    }
                        //折线:双旋
                        //    g
                        //  p   u
                        //    c
                    else{
                        RotateL(parent);
                        RotateR(grandfather);
                        cur->_col = BLACK;
                        grandfather->_col = RED;
                    }
                    //旋转后的根均为黑色，可以跳出
                    break;
                }
            }
                //p在g的右
                //grandfather->_right == parent;
            else{
                Node* uncle = grandfather->_left;
                //情况1：u存在且为红(变色处理，并继续往上处理)
                if (uncle && uncle->_col == RED){
                    parent->_col = BLACK;
                    uncle->_col = BLACK;
                    grandfather->_col = RED;

                    //继续向上调整
                    cur = grandfather;
                    parent = cur->_parent;
                }
                    //情况2：u不存在/u存在且为黑(旋转+变色)
                else{
                    //斜线:单旋
                    //    g
                    //  u   p
                    //        c
                    if (cur == parent->_right){
                        RotateL(grandfather);
                        parent->_col = BLACK;
                        grandfather->_col = RED;
                    }
                        //折线:双旋
                        //    g
                        //  u   p
                        //    c
                    else{
                        RotateR(parent);
                        RotateL(grandfather);
                        cur->_col = BLACK;
                        grandfather->_col = RED;
                    }
                    //旋转后的根均为黑色，可以跳出
                    break;
                }
            }
        }

        _root->_col = BLACK;
        return make_pair(iterator(newNode), true);
    }

    void InOrder(){
        _InOrder(_root);
        cout << endl;
    }

    bool IsBalance(){
        //检查根节点颜色
        if (_root && _root->_col == RED){
            cout << "The node's colour of root is red" << endl;
            return false;
        }

        //计算路径上黑色节点个数的基准值
        int benchMark = 0;
        Node* cur = _root;
        while (cur){
            if (cur->_col == BLACK)
                ++benchMark;
            cur = cur->_left;
        }

        //判断红色节点是否连续和黑色节点数量是否相等
        return _Check(_root,0,benchMark);
    }

    int Height(){
        return _Height(_root);
    }

private:
    void RotateL(Node* parent){
        Node* subR = parent->_right;
        Node* subRL = subR->_left;

        parent->_right = subRL;
        if (subRL)
            subRL->_parent = parent;

        Node* pParent = parent->_parent;

        subR->_left = parent;
        parent->_parent = subR;

        if (pParent == nullptr){
            _root = subR;
            subR->_parent = nullptr;
        } else{
            if (pParent->_right == parent)
                pParent->_right = subR;
            else
                pParent->_left = subR;

            subR->_parent = pParent;
        }
    }

    void RotateR(Node* parent){
        Node* subL = parent->_left;
        Node* subLR = subL->_right;

        parent->_left = subLR;
        if (subLR)
            subLR->_parent = parent;

        Node* pParent = parent->_parent;

        subL->_right = parent;
        parent->_parent = subL;

        if (parent == _root){
            _root = subL;
            subL->_parent = nullptr;
        } else{
            if (pParent->_right == parent)
                pParent->_right = subL;
            else
                pParent->_left = subL;
            subL->_parent = pParent;
        }
    }

    void _InOrder(Node* root){
        KeyOfT kot;
        if (root == nullptr)
            return;
        _InOrder(root->_left);
        cout << kot(root->_data) << " ";
        _InOrder(root->_right);
    }

    void _Destroy(Node* root){
        if (root == nullptr)
            return;
        _Destroy(root->_left);
        _Destroy(root->_right);
        delete root;
    }

    int _Height(Node* root){
        if (root == nullptr)
            return 0;
        int leftHeight = _Height(root->_left);
        int rightHeight = _Height(root->_right);

        return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
    }

    bool _Check(Node* root, int blackNum, int benchMark){
        if (root == nullptr){
            if (blackNum != benchMark){
                cout << "There exist unequal numbers of black nodes" << endl;
                return false;
            }

            return true;
        }

        if (root->_col == BLACK)
            ++blackNum;

        if (root->_col == RED && root->_parent->_col == RED){
            cout << "There exist continuous red node" << endl;
            return false;
        }

        return _Check(root->_left, blackNum, benchMark) && _Check(root->_right, blackNum, benchMark);
    }

    Node* _root = nullptr;
};
